﻿#include "infer_cls_v2.h"
#include "common/define.h"
#include <opencv2/opencv.hpp>

namespace txr_algo_dlm_cls
{ 



	//===================================================================================
	InferClsV2::InferClsV2()
	{
	}
	InferClsV2::~InferClsV2()
	{
	}

	bool InferClsV2::Init(int gpu_id, const char* dats_path)
	{
		if (!InferBase::Init(gpu_id,dats_path))
		{
			return false;
		}
		//......
		return true;
	}

	bool InferClsV2::LoadEngine(std::vector<char> v_engine_data)
	{
		if (!InferBase::LoadEngine(v_engine_data))
		{
			return false;
		}
		if (m_engine->getNbBindings() != 2)
		{
			return true;
		}
		assert(m_engine->getNbBindings() == 2);
		int nbBindings = m_engine->getNbBindings();
		m_v_dev_buffer.resize(2);
		m_v_dev_buf_size.resize(2);
		
		for (int i = 0; i < nbBindings; ++i) {
			nvinfer1::Dims dims = m_engine->getBindingDimensions(i);
			nvinfer1::DataType dtype = m_engine->getBindingDataType(i);
			int64_t totalSize = volume(dims) * 1 * getElementSize(dtype);
			m_v_dev_buf_size[i] = totalSize;
			std::cout << "binding" << i << ": " << totalSize << std::endl;
		}
		for (int i = 0; i < m_v_dev_buffer.size(); ++i)
		{
			cudaMalloc(&m_v_dev_buffer[i], m_v_dev_buf_size[i]);
		}
		m_out_size = m_v_dev_buf_size[1] / sizeof(float) / m_info.batch_size;

		return true;
	}
	void RecordCuImage(void * dev_ptr,int w,int h)
	{
		static int test_count = 0;
		{
			using namespace cv;
			std::vector<float> v_img;
			v_img.resize(w * h * 3);
			cudaMemcpy(v_img.data(), dev_ptr, w * h * 3 * sizeof(float), cudaMemcpyDeviceToHost);
			Mat mat;
			mat.create(h, w, CV_8UC3);
			int img_size = mat.rows * mat.cols;
			int index = 0;
			for (int i = 0; i < mat.rows; ++i)
			{
				for (int j = 0; j < mat.cols; ++j)
				{
					Vec3b rgb;
					rgb[2] = qBound(0.f,(v_img[index]) * 255.f,255.f);
					rgb[1] = qBound(0.f,(v_img[index + img_size]) * 255.f, 255.f);
					rgb[0] = qBound(0.f,(v_img[index + img_size + img_size]) * 255.f, 255.f);
					mat.at<Vec3b>(i, j) = rgb;
					index++;
				}
			}
			test_count++;
			imwrite("resize_" + std::to_string(test_count) + ".png", mat);
		}
	}

	void InferClsV2::Detect(st_detect_unit* p_unit, int num)
	{
		int batch_size = m_info.batch_size;
		if (num > batch_size)
		{
			std::cout << "input unit size > batch size!" << std::endl;
			num = batch_size;
		}

		auto a_start = std::chrono::high_resolution_clock::now();

		int net_w = m_info.image_width,
			net_h = m_info.image_height,
			net_c = m_info.input_channel;

		if (_v_resize_space.size() < batch_size)
		{
			for (auto i : _v_resize_space)
			{
				CuImgResizeFree(i);
			}
			_v_resize_space.resize(batch_size);
		}
		float* dlm_input_img = (float*)m_v_dev_buffer[0];
		cudaMemset(m_v_dev_buffer[0], m_v_dev_buf_size[0],0);
		int input_ptr_index = 0;
		int resize_h = 0;
		int resize_w = 0;
		{
			float mean[3] = { m_info.img_mean[0], m_info.img_mean[1], m_info.img_mean[2] };
			float dev[3] = { m_info.img_std[0], m_info.img_std[1], m_info.img_std[2] };
			for (int i = 0; i < batch_size; ++i)
			{
				//根据输入的图像数对应图像数据，如果输入数小于batch用第一个图像数据填充
				st_detect_unit& data = i < num ? p_unit[i] : p_unit[0];
				st_img_rgb& img = data.img;
	// 			std::swap(img.pp_rgbf[0], img.pp_rgbf[2]);
	// 			std::swap(img.pp_rgb8[0], img.pp_rgb8[2]);
	// 			std::swap(img.pp_rgb16[0], img.pp_rgb16[2]);

				float ratio = float(img.w) / float(img.h);
				int resize_w, resize_h;
				if (ceilf(m_info.image_height * ratio) > m_info.image_width)
					resize_w = m_info.image_width;
				else
					resize_w = int(ceilf(m_info.image_height * ratio));
				resize_h = m_info.image_height;
	
				st_cuda_resize_dev_space& space = _v_resize_space[i];
	
				int max_line_size = __max(__max(img.w, net_w), __max(img.h, net_h));
				CuImgResizeMalloc(space, max_line_size * max_line_size);
	
				space.img_rgbf_net = dlm_input_img + input_ptr_index;
				if (img.pp_rgb8[0])
				{
					CuImgResizeRgb8(img.pp_rgb8, img.w, img.h, resize_w, resize_h, space, net_w, net_h,   mean, dev);
				}
				else if (img.pp_rgb16[0])
				{
					CuImgResizeRgb16(img.pp_rgb16, img.w, img.h, resize_w, resize_h, space, net_w, net_h, mean, dev);
				}
				else if (img.pp_rgbf[0])
				{
					CuImgResizeRgbF(img.pp_rgbf, img.w, img.h, resize_w, resize_h, space, net_w, net_h,  mean, dev);
				}
				//RecordCuImage(dlm_input_img+ input_ptr_index, net_w, net_h);
				input_ptr_index += net_w * net_h * net_c;
			}
		}

		// do inference
		{
			auto t_start = std::chrono::high_resolution_clock::now();
			m_context->executeV2(m_v_dev_buffer.data());
			auto t_end = std::chrono::high_resolution_clock::now();
			float total_inf = std::chrono::duration<float, std::milli>(t_end - t_start).count();
			std::cout << "Inference take: " << total_inf << " ms." << std::endl;
		}

		//postProcess
		{
			auto p_start = std::chrono::high_resolution_clock::now();
			
			int out_buf_size = m_out_size * m_info.batch_size* sizeof(float);
			if (m_v_out_buf.size() < out_buf_size)
			{
				m_v_out_buf.resize(out_buf_size);
			}
			cudaMemcpy(m_v_out_buf.data(), m_v_dev_buffer[1], out_buf_size, cudaMemcpyDeviceToHost);
			postProcess(p_unit, num, m_v_out_buf.data(), m_out_size);
			auto p_end = std::chrono::high_resolution_clock::now();
			float total_inf = std::chrono::duration<float, std::milli>(p_end - p_start).count();
			std::cout << "postProcess take: " << total_inf << " ms." << std::endl;
		}
	}

	void InferClsV2::postProcess(st_detect_unit* p_unit, int size, float* output, const int& out_size)
	{
		int index = 0;
		for (auto i = 0; i < size; ++i)
		{
			float* out = output + index * out_size;
			auto max_pos = std::max_element(out, out + out_size);

			p_unit[i].result.cls = max_pos - out;
			p_unit[i].result.prob = out[p_unit[i].result.cls];
			index++;
		}
	}

}//txr_algo_dlm_cls